import os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from skimage import io
from torch import optim
from torch.utils.data import DataLoader
from torchvision import transforms

from model import MaxVitUnet


class Dataset(torch.utils.data.Dataset):
    def __init__(self, root_dir, target_size=256, min_size=64, max_size=1024, is_train=True):
        self.root_dir = root_dir
        self.target_size = target_size
        self.min_size = min_size
        self.max_size = max_size
        self.is_train = is_train

        self.folders = [d for d in os.listdir(root_dir)
                        if os.path.isdir(os.path.join(root_dir, d))]

    def process_image_and_mask(self, image, mask):
        # 1. 检查图片尺寸是否在合理范围内
        h, w = image.shape[-2:]
        if h < self.min_size or w < self.min_size:
            scale = self.min_size / min(h, w)
            new_h, new_w = int(h * scale), int(w * scale)
            image = F.interpolate(image.unsqueeze(0), size=(new_h, new_w),
                                  mode='bilinear', align_corners=False).squeeze(0)
            mask = F.interpolate(mask.unsqueeze(0), size=(new_h, new_w),
                                 mode='nearest').squeeze(0)

        elif h > self.max_size or w > self.max_size:
            scale = self.max_size / max(h, w)
            new_h, new_w = int(h * scale), int(w * scale)
            image = F.interpolate(image.unsqueeze(0), size=(new_h, new_w),
                                  mode='bilinear', align_corners=False).squeeze(0)
            mask = F.interpolate(mask.unsqueeze(0), size=(new_h, new_w),
                                 mode='nearest').squeeze(0)

        # 2. 保持长宽比的调整策略
        h, w = image.shape[-2:]
        if h > w:
            new_h = self.target_size
            new_w = int(w * (new_h / h))
        else:
            new_w = self.target_size
            new_h = int(h * (new_w / w))

        image = F.interpolate(image.unsqueeze(0), size=(new_h, new_w),
                              mode='bilinear', align_corners=False).squeeze(0)
        if mask is not None:
            mask = F.interpolate(mask.unsqueeze(0), size=(new_h, new_w),
                                 mode='nearest').squeeze(0)

        # 3. 填充到目标尺寸
        pad_h = self.target_size - new_h
        pad_w = self.target_size - new_w

        pad_h1, pad_h2 = pad_h // 2, pad_h - (pad_h // 2)
        pad_w1, pad_w2 = pad_w // 2, pad_w - (pad_w // 2)

        padding = (pad_w1, pad_w2, pad_h1, pad_h2)
        image = F.pad(image, padding, mode='constant', value=0)
        if mask is not None:
            mask = F.pad(mask, padding, mode='constant', value=0)

        return image, mask

    def __getitem__(self, idx):
        folder_name = self.folders[idx]
        folder_path = os.path.join(self.root_dir, folder_name)

        # Load image
        image_path = os.path.join(folder_path, 'images', os.listdir(os.path.join(folder_path, 'images'))[0])
        image = io.imread(image_path)

        if image.shape[-1] == 4:
            image = image[..., :3]
        image = np.transpose(image, (2, 0, 1))
        image = torch.from_numpy(image).float() / 255.0

        if self.is_train:
            # Load and combine masks with explicit normalization
            mask_dir = os.path.join(folder_path, 'masks')
            mask_files = sorted(os.listdir(mask_dir))
            combined_mask = None

            for mask_file in mask_files:
                mask_path = os.path.join(mask_dir, mask_file)
                # 确保mask是灰度图并归一化到0-1
                mask = io.imread(mask_path, as_gray=True)
                # 如果mask不是0-1范围，进行归一化
                if mask.max() > 1:
                    mask = mask / 255.0
                if combined_mask is None:
                    combined_mask = mask
                else:
                    combined_mask = np.clip(combined_mask + mask, 0, 1)

            # 确保mask是浮点数类型并在0-1范围内
            combined_mask = combined_mask.astype(np.float32)
            combined_mask = np.clip(combined_mask, 0, 1)
            combined_mask = torch.from_numpy(combined_mask).float()
            combined_mask = combined_mask.unsqueeze(0)  # Add channel dim

            # Process image and mask
            image, mask = self.process_image_and_mask(image, combined_mask)

            # 最后再次确保mask在0-1范围内
            mask = torch.clamp(mask, 0, 1)

            return image, mask, {'folder_name': folder_name}
        else:
            # Process image only
            image, _ = self.process_image_and_mask(image, None)
            return image, {'folder_name': folder_name}

    def __len__(self):
        return len(self.folders)


from tqdm import tqdm


def train(model, dataloader, criterion, optimizer, device, num_epochs=31):
    model.train()
    for epoch in range(num_epochs):
        print(f'第{epoch + 1}轮训练:')
        running_loss = 0.0  # 用于计算每个epoch的平均损失
        # 在数据加载器上添加进度条
        for images, masks, _ in tqdm(dataloader, desc=f"Epoch {epoch + 1}/{num_epochs}", leave=False):
            images = images.float().to(device)
            masks = masks.float().to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * images.size(0)  # 累加损失

        epoch_loss = running_loss / len(dataloader.dataset)
        print(f'Epoch {epoch + 1}, Loss: {epoch_loss:.4f}')
        if (epoch + 1) % 3 == 0:
            torch.save(model.state_dict(), f'MaxU-{epoch}.pt')


if __name__ == '__main__':
    dataset = Dataset(
        root_dir='/Volumes/For_Mac/unet++/unet++/inputs/train',
        target_size=256,
        min_size=64,
        max_size=1024
    )

    # 数据集与数据加载器
    # 可以根据需求进行数据增强操作
    transform = transforms.Compose([
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

    # 初始化模型、损失函数和优化器
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = MaxVitUnet(3, 1).to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.BCELoss()
    # 开始训练
    train(model, dataloader, criterion, optimizer, device)
